﻿#include "pch.h"
#include <stdio.h>
#include <Windows.h>
#include <memoryapi.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>

struct PatchDef {
    const char *        name;
    uint32_t *          offsets;
    int                 size;
};

#define ARRAY_SIZE(a)   (sizeof(a) / sizeof(a[0]))
#define _INIT(...)                                                      \
    .offsets = (uint32_t *)&(uint32_t []){ __VA_ARGS__},                \
    .size = sizeof((uint32_t[]){ __VA_ARGS__}) / sizeof(uint32_t)


static const unsigned char      expectedPattern[] = { 0xff, 0x15 };
struct PatchDef                 patches[] = {
    {
        .name           = "canjiren",
        _INIT(80776, 80894, 80947),
    },
    {
        .name           = "feng1",
        _INIT(149551, 149669, 149722),
    },
    {
        .name           = "feng2",
        _INIT(170323, 170441, 170494)
    },
};


static void loadMap();
static void runPatch(const char *name);
static void trace(const char *fmt, ...);

BOOL APIENTRY DllMain( HMODULE hModule,
                       DWORD  ul_reason_for_call,
                       LPVOID lpReserved
                     )
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
        loadMap();
        break;
    case DLL_THREAD_ATTACH:
        trace("Attached\n");
        break;
    case DLL_THREAD_DETACH:
    case DLL_PROCESS_DETACH:
        trace("Detached\n");
        break;
    }
    return TRUE;
}

//#define D2F_DEBUG

static void trace(const char *fmt, ...)
{
#ifdef D2F_DEBUG
    static FILE *fp;

    va_list ap;
    if (fp == NULL) {
        fopen_s(&fp, "d2f.log", "a+");
        if (fp == NULL) {
            return;
        }
    }
    va_start(ap, fmt);
    vfprintf(fp, fmt, ap);
    va_end(ap);
    fflush(fp);
    fclose(fp);
    fp = NULL;
#endif
}

static WNDPROC getOrigProc(HWND hwnd)
{
    TCHAR       className[64];
    int         classNameLen;
    WNDCLASS    classInfo;
    HINSTANCE   hInst = GetModuleHandle(NULL);
    WNDPROC     orig;

    if (hInst == NULL) {
        return NULL;
    }
    // Get "Diablo II" class name
    classNameLen = GetClassName(hwnd, &className[0], sizeof(className) / sizeof(className[0]));
    if (classNameLen == 0) {
        trace("Failed to get class name\n");
        return FALSE;
    }
    if (!GetClassInfo(hInst, className, &classInfo)) {
        trace("Failed to get class info\n");
        return FALSE;
    }
    orig = classInfo.lpfnWndProc;
    trace("Orig proc: %p\n", orig);
    return orig;
}

__declspec(noinline)
LRESULT WINAPI mySendMessage(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
    static WNDPROC      origProc;
    static BOOL         isGot;

    LRESULT             ret;

    trace("Get message 0x%x param 0x%x 0x%x\n", msg, wParam, lParam);
    if (!isGot) {
        origProc = getOrigProc(hWnd);
        isGot = TRUE;
    }
    if (origProc != NULL) {
        ret = CallWindowProcA(origProc, hWnd, msg, wParam, lParam);
    } else {
        ret = SendMessageA(hWnd, msg, wParam, lParam);
    }

    return ret;
}


static void makeCall(unsigned char *buf, int size, void * old, void * func)
{
    uint32_t    oldAddr, newAddr;
    int         offset;

    oldAddr = (uint32_t)(uintptr_t)old;
    newAddr = (uint32_t)(uintptr_t)func;
    buf[0] = 0xE8;
    *(uint32_t *)&buf[1] = (newAddr - (oldAddr + 5));
    offset = 5;
    for (; offset < size; offset += 1) {
        buf[offset] = 0x90;
    }
}

static bool readLocalBytes(void *addr, unsigned char *buf, int len)
{
    DWORD       oldProtect;
    bool        ok;

    memset(buf, 0, len);
    if (!VirtualProtect(addr, len, PAGE_EXECUTE_READWRITE, &oldProtect)) {
        trace("Failed to set protection addr %p\n", addr);
        return false;
    }
    ok = ReadProcessMemory(GetCurrentProcess(), addr, (void *)buf, len, NULL);
    VirtualProtect(addr, len, oldProtect, &oldProtect);
    return ok;
}

static bool writeLocalBytes(void *addr, const unsigned char *buf, int len,
                           const unsigned char *expected, int checkLen)
{
    DWORD       oldProtect;
    uint32_t    data[2];
    bool        ok;

    if (!VirtualProtect(addr, len, PAGE_EXECUTE_READWRITE, &oldProtect)) {
        return false;
    }
    ReadProcessMemory(GetCurrentProcess(), addr, (void *)data, sizeof(data), NULL);

    ok = false;
    if (expected != NULL) {
        unsigned char *         orig = (unsigned char *)data;
        for (int i = 0; i < checkLen; i += 1) {
            if (expected[i] != orig[i]) {
                trace("Data mismatch offset %u (0x%x 0x%x)\n",
                      (uintptr_t)addr + i, orig[i], expected[i]);
                goto out;
            }
        }
    }
    trace("Patch data at %p, orig %08x %08x\n", addr, data[0], data[1]);
    ok = WriteProcessMemory(GetCurrentProcess(), addr, buf, len, NULL);

out:
    VirtualProtect((void *)addr, len, oldProtect, &oldProtect);
    return ok;
}

static void patch(void *addr, void *func, int len)
{
    unsigned char       buf[12];

    makeCall(buf, len, addr, func);
    writeLocalBytes(addr, buf, len, expectedPattern, sizeof(expectedPattern));
}

static void patchAtOffset(uint32_t base, uint32_t offset)
{
    patch((void *)(uintptr_t)(base + offset), mySendMessage, 6);
}

static bool isPatchValid(uintptr_t base, const struct PatchDef *patch)
{
    unsigned char       buf[sizeof(expectedPattern)];

    for (int i = 0; i < patch->size; i += 1) {
        void *          addr = (void *)(base + patch->offsets[i]);
        if (!readLocalBytes(addr, &buf[0], sizeof(buf))) {
            return false;
        }
        if (memcmp(buf, expectedPattern, sizeof(expectedPattern)) != 0) {
            return false;
        }
    }
    return true;
}

__declspec(noinline)
static struct PatchDef * getValidPatch(uintptr_t base)
{
    struct PatchDef *   selected = NULL;
    for (int i = 0; i < ARRAY_SIZE(patches); i += 1) {
        if (isPatchValid(base, &patches[i])) {
            selected = &patches[i];
            break;
        }
    }
    if (selected == NULL) {
        trace("No valid patch found\n");
        return NULL;
    }
    trace("Selected %s for patching\n", selected->name);
    trace("Offsets:");
    for (int i = 0; i < selected->size; i += 1) {
        trace(" %u", selected->offsets[i]);
    }
    trace("\n");

    return selected;
}

static bool getDllPath(char *path, int size)
{
    HMODULE     hm;
    if (!GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
                           GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
                           (LPCSTR)&getValidPatch, &hm)) {
        trace("GetModuleHandle failed %d\n", GetLastError());
        return false;
    }
    if (!GetModuleFileNameA(hm, path, size)) {
        trace("Failed to get file name\n");
        return false;
    }

    return true;
}

static bool tryLoadMap(const char *name)
{
    char                path[512];
    char                dllPath[512];
    HMODULE             module;
    char *              div;

    if (!getDllPath(path, sizeof(path))) {
        return false;
    }
    trace("Path: %s\n", path);
    div = strrchr(path, '\\');
    if (div == NULL) {
        div = path;
    }
    *div = '\0';
    snprintf(dllPath, sizeof(dllPath), "%s\\%s", path, name);
    trace("Load dll %s\n", dllPath);

    module = LoadLibraryA(dllPath);
    if (module == NULL) {
        trace("Unload to load map %s\n", dllPath);
        return false;
    }
    return true;
}

static void loadMap()
{
    const char *        origDllName = "_d2hackmap.dll";
    const char *        defaultDllName = "d2hackmap.dll";
    const char *        name;

    if (tryLoadMap(origDllName)) {
        name = origDllName;
        // Prevent me from unloading
        LoadLibraryA(defaultDllName);
    } else {
        name = defaultDllName;
    }
    runPatch(name);
}

static void runPatch(const char *name)
{
    HMODULE             module;
    uint32_t            base;
    struct PatchDef *   patch;

    module = GetModuleHandleA(name);
    if (module == NULL) {
        trace("Run patch failed, cannot not find module %s\n", name);
        return;
    }
    base = (uintptr_t)module;
    patch = getValidPatch(base);
    if (patch == NULL) {
        return;
    }
    trace("Dll base %p, patch function %p\n", module, mySendMessage);
    for (int i = 0; i < patch->size ; i += 1) {
        patchAtOffset(base, patch->offsets[i]);
    }
}
